package com.pine.web.core;

import com.pine.common.beans.exception.BusinessException;
import com.pine.common.beans.exception.ExceptionCodeEnum;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.lang.NonNull;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

/**
 * 限流拦截器
 *
 * @author pine manage
 * @since 2024-08-14

 */
@Slf4j
public class RateLimitInterceptor implements HandlerInterceptor {

    private final RRateLimiter rateLimiter;

    public RateLimitInterceptor(RedissonClient redisson) {
        this.rateLimiter = redisson.getRateLimiter("apiRateLimiter");
        /*
         * RateType.OVERALL：表示限流规则应用于所有客户端。这意味着不论有多少客户端进行请求，总的请求频率都受到这个规则的限制
         * RateType.PER_CLIENT：表示限流规则应用于每个客户端。这意味着每个客户端的请求频率都受到这个规则的限制
         * 10：表示每秒允许的最大请求数量。在这个例子中，每秒最多允许 10 个请求。
         * 1：表示时间窗口的长度，这里是 1 秒。
         * RateIntervalUnit.SECONDS：表示时间窗口的单位，这里是秒。
         */
        rateLimiter.trySetRate(RateType.OVERALL, 10, 1, RateIntervalUnit.SECONDS);
    }

    @Override
    public boolean preHandle(@NonNull HttpServletRequest request, @NonNull HttpServletResponse response, @NonNull Object handler) throws Exception {
        // 如果请求的是方法，则需要做校验
        if (handler instanceof HandlerMethod handlerMethod) {
            // 获取目标方法上是否有指定注解
            RateLimit rateLimit = handlerMethod.getMethodAnnotation(RateLimit.class);
            if (rateLimit == null) {
                // 如果没有注解，则直接放行
                return true;
            }

            // 尝试获取许可
            boolean permit = rateLimiter.tryAcquire(1);

            if (!permit) {
                // 超出限流，返回错误响应
                response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
                throw new BusinessException(ExceptionCodeEnum.TOO_MANY_REQUESTS);
            }
        }

        // 允许请求继续
        return true;
    }
}
